{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  Ochre.Compiler.Alpha
-- Copyright   :  (c) Stephen Tetley 2012
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Alpha renaming for K Normal Form.
--
--------------------------------------------------------------------------------

module Ochre.Compiler.Alpha
  (
    Env
  , find
  , findVal
  , inserts  
  
  , alpha

  ) where


import Ochre.Compiler.CompilerMon
import Ochre.Compiler.Syntax
import Ochre.Compiler.VarId

import Control.Applicative
import qualified Data.Map as M


--------------------------------------------------------------------------------
-- alpha renaming

type Env = M.Map VarId VarId


-- | We need the permissiveness of returning original VarId to 
-- handle variables defined with vardef that wont be renamed
--
find :: Env -> VarId -> Compiler VarId
find env x = maybe (return x) return $ M.lookup x env


findVal :: Env -> Value -> Compiler Value
findVal env (Var v)     = Var <$> find env v
findVal env (Signal v)  = Signal <$> find env v
findVal _   v           = pure v

inserts :: [(VarId,VarId)] -> Env -> Env
inserts vvs env = foldr (\(a,b) acc -> M.insert a b acc) env vvs


-- | Re-initialize counts to zero...
--
alpha :: Expr -> Compiler Expr
alpha e = initState >> ac M.empty e

ac :: Env -> Expr -> Compiler Expr
ac env (Return val)                 = Return <$> acV env val

-- Don\'t alpha rename opcodes.
ac env (OpCall pid vs)              = OpCall pid <$> mapM (acV env) vs

ac env (LetValue vid val e1)        = do
    vx <- fresh (varRate vid)
    va <- acV env val
    ea <- ac (M.insert vid vx env) e1
    return $ LetValue vx va ea

ac env (LetSig1 vid e1 e2)          = do
    vx <- fresh (varRate vid)
    ea <- ac env e1
    eb <- ac (M.insert vid vx env) e2
    return $ LetSig1 vx ea eb

ac env (LetSig2 v1 v2 e1 e2)        = do
    xa <- fresh (varRate v1)
    xb <- fresh (varRate v2)
    ea <- ac env e1
    eb <- ac (inserts [(v1,xa), (v2,xb)] env) e2
    return $ LetSig2 xa xb ea eb

ac env (LetSig3 v1 v2 v3 e1 e2)     = do
    xa <- fresh (varRate v1)
    xb <- fresh (varRate v2)
    xc <- fresh (varRate v3)
    ea <- ac env e1
    eb <- ac (inserts [(v1,xa), (v2,xb), (v3,xc)] env) e2
    return $ LetSig3 xa xb xc ea eb

ac env (LetSig4 v1 v2 v3 v4 e1 e2)  = do
    xa <- fresh (varRate v1)
    xb <- fresh (varRate v2)
    xc <- fresh (varRate v3)
    xd <- fresh (varRate v4)
    ea <- ac env e1
    eb <- ac (inserts [(v1,xa), (v2,xb), (v3,xc), (v4,xd)] env) e2
    return $ LetSig4 xa xb xc xd ea eb

ac env (Ifte v1 e1 e2)              = 
    Ifte <$> findVal env v1 <*> ac env e1 <*> ac env e2

ac env (WriteSig vid e)             = WriteSig <$> find env vid <*> ac env e

ac env (e1 :>> e2)                  = 
    (:>>) <$> ac env e1 <*> ac env e2



acV :: Env -> Value -> Compiler Value
acV env (Var vid)               = Var <$> find env vid
acV env (Signal vid)            = Signal <$> find env vid
acV env (FunCallV v vs)         = FunCallV v <$> mapM (acV env) vs
acV env (UnaryV op v)           = UnaryV op <$> acV env v
acV env (BinV op v1 v2)         = BinV op <$> acV env v1 <*> acV env v2
acV env (RelV op v1 v2)         = RelV op <$> acV env v1 <*> acV env v2
acV env (CondV v1 vt vf)        = 
    CondV <$> acV env v1 <*> acV env vt <*> acV env vf

acV _   val                     = return val
